/** @file download.cpp
    @brief Image downloader */
#include "download.h"

#define DUMMY_IMAGE_EXTENSION	".qrv"

Downloader::Downloader() {
	m_curl = NULL;
}

bool Downloader::init() {
	m_curl = curl_easy_init();

	srand( time( NULL ) );

	if (m_curl)
		return true;

	printf( "E: Failed to initialize libCURL..\n" );
	return false;
}

Downloader::~Downloader() {
	if (m_curl)
		curl_easy_cleanup( m_curl );
}

std::string Downloader::query() {
	CURLcode result;
	FILE *fo = NULL;
	std::ostringstream stm;

	if (!m_curl)
		return "";

	// Pick a random number
	stm << rand() % 1000;

	std::string base = "http://www.google.com/images?q=";
	std::string query = base + stm.str();
	std::string dest = "pics/0.html";

	fo = fopen( dest.c_str(), "wt" );

	if (!fo)
		return "";

	// Set the URL of the query
	curl_easy_setopt( m_curl, CURLOPT_URL, query.c_str() );
    // Set destination write function
    curl_easy_setopt( m_curl, CURLOPT_WRITEFUNCTION, fwrite );
   	// Set destination file
	curl_easy_setopt( m_curl, CURLOPT_WRITEDATA, fo );
	// Set verbose mode
	curl_easy_setopt( m_curl, CURLOPT_VERBOSE, 1 );
	// Perform the query
	result = curl_easy_perform( m_curl );

	fclose( fo );

	if (result == 0)
		return dest;
	return "";
}

std::vector< std::string > Downloader::parse( const std::string &s_file ) {
	std::ifstream ifs( s_file.c_str() );
	std::vector< std::string > result;
	std::string line, url;
	size_t start = 0, end = 0;

	while ( !ifs.eof() ) {
		start = 0;
		end = -1;

		std::getline( ifs, line );

        // Find embedded images
		while ( true ) {
			// Find strings starting with "http://"
			start = line.find( "http://", end + 1 );
			
			if (start == std::string::npos)
				break;

			// Find the end of the image source (" or \)
			end = line.find_first_of( "\"\\", start + 8 );

			if (end == std::string::npos)
				break;

			// Store the URL
			url = line.substr( start, end - start );
			// Validate it
			url = validate_url( url );
			// And add it to the result if not empty
			if (!url.empty()) {
				result.push_back( url );
            }
		}
	}

	return result;
}

std::string Downloader::validate_regular_image( const std::string &s_url ) {
	size_t pos = s_url.find( "http://", 0 );
	// All of the URLs should start with "http://"
	if (pos == std::string::npos)
		return "";

	// Make sure it's not a .svg.png file (which can't be opened by QT)
	size_t dot = s_url.find( ".svg.png" );
	// Ignore this kind of file extensions
	if (dot != std::string::npos)
		return "";

	// Check if it's a valid picture file
	dot = s_url.find_last_of( "." );
	std::string temp = s_url.substr( dot, std::string::npos );
	// Only allow for jpg, png files (QT can't display gif-s)
	if (temp == ".jpg" || temp == ".png")
		return s_url.substr( pos, std::string::npos );
	return "";
}

std::string Downloader::validate_gstatic_image( const std::string &s_url ) {
	size_t pos = s_url.find( "http://", 0 );
	// All of the URLs should start with "http://"
	if (pos == std::string::npos)
		return "";

	size_t gstatic_pos = s_url.find( ".gstatic.com/images?q", 8 );
	// All of the URLs should have the form "http://*.gstatic.com/images?q*"
	if (gstatic_pos == std::string::npos)
		return "";

	return s_url.substr( pos, std::string::npos );
}

std::string Downloader::validate_url( const std::string &s_url ) {
	printf( ">> \"%s\"\r\n", s_url.c_str() );
	
	std::string img_url = "";
	
	// Is it a regular image?
	img_url = validate_regular_image( s_url );
	if (img_url != "")
		return img_url;
	
	// Is it a google static image?
	img_url = validate_gstatic_image( s_url );
	if (img_url != "")
		return img_url;

	return img_url;
}

std::string Downloader::get_dest( const std::string &s_url ) {
	static int num_pics = 0;
	std::string result = "pics/";
	std::ostringstream stm;
	
	std::string img_url = "";
	
	stm << num_pics;
	
	// Is it a regular image?
	img_url = validate_regular_image( s_url );
	if (img_url != "") {
		// Find the extension
		size_t dot = s_url.find_last_of( "." );
		if (dot == std::string::npos)
			return "";
		// Generate a file name
		result += stm.str() + s_url.substr( dot, std::string::npos );
	}
	
	// Is it a google static image?
	img_url = validate_gstatic_image( s_url );
	if (img_url != "") {
		// Generate a file name with a dummy extension
		result += stm.str() + DUMMY_IMAGE_EXTENSION;
	}

	num_pics++;

	return result;
}

std::string Downloader::download( const std::string &s_image ) {
	CURLcode result;
	FILE *fo = NULL;

	if (!m_curl)
		return "";

	// Find the destination file
	std::string dest = get_dest( s_image );
	
	printf( "I: Downloading image \"%s\" and saving it as \"%s\"\n", s_image.c_str(), dest.c_str() );
	
	// Open it
	fo = fopen( dest.c_str(), "wb" );

	if (!fo)
		return "";

	// Set the URL of the query
	curl_easy_setopt( m_curl, CURLOPT_URL, s_image.c_str() );
	// Set destination file
	curl_easy_setopt( m_curl, CURLOPT_WRITEDATA, fo );
	// Perform the query
	result = curl_easy_perform( m_curl );

	fclose( fo );
	return dest;
}

std::string Downloader::validate_image_file( const std::string &s_path) {
	// Find the extension
	size_t dot = s_path.find_last_of( "." );
	if (dot == std::string::npos)
		return "";
	// Extract file path before the extension
	std::string pre_ext = s_path.substr( 0, dot );
	// Extract file extension
	std::string ext = s_path.substr( dot, std::string::npos );
	// New path for the file
	std::string new_path = s_path;
	
	// The file needs renaming based on file format?
	if (ext == DUMMY_IMAGE_EXTENSION) {
		unsigned char file_header [16];
		
		FILE *fi = fopen( s_path.c_str(), "rb" );
		
		if (!fi)
			return "";
		
		if (fread( file_header, 1, sizeof( file_header ), fi ) != sizeof( file_header )) {
			printf( "E: Failed to read the header of \"%s\"\r\n", s_path.c_str() );
			fclose( fi );
			return "";
		}
		
		fclose( fi );
		
		// It's a PNG?
		if (file_header [0] == 0x89 && file_header [1] == 'P' && file_header [2] == 'N' && file_header [3] == 'G') {
			new_path = pre_ext + ".png";
			rename( s_path.c_str(), new_path.c_str() );
		} else if (file_header [0] == 0xFF && file_header [1] == 0xD8 && file_header [2] == 0xFF) {
			new_path = pre_ext + ".jpg";
			rename( s_path.c_str(), new_path.c_str() );
		} else if (file_header [0] == 'G' && file_header [1] == 'I' && file_header [2] == 'F') {
			new_path = pre_ext + ".gif";
			rename( s_path.c_str(), new_path.c_str() );
		}
	}
	
	return new_path;
}

std::string Downloader::fetch() {
	std::vector< std::string > images;
	std::string result;
	int id = 0;

    // Not initialized?
	if (!m_curl)
		init();

	if (m_curl) {
        // Timeout queries after 10 attempts
		for (int i=0; i<10; i++) {
			if (!result.empty())
				break;

			result = query();
		}
        // No luck?
		if (result.empty()) {
			printf( "E: Failed the query 10 times, aborting..\n" );
			return "";
		}

        // Parse the query result
		images = parse( result );

        // No images?
        if (images.empty()) {
            printf( "E: Parsing failed (couldn't find any images)..\n" );
            return "";
        }

        // Dump a list of all found images
		for (int i=0; i<(int)images.size(); i++)
			printf( "IMG \"%s\"\n", images [i].c_str() );

        // Pick a random image
		id = rand() % (images.size() - 1);
		// Download it
		result = download( images [id] );
		// Check the format and rename it, if necessary
		result = validate_image_file( result );

		printf( "Done..\n" );
	}

    // Return the filename of the downloaded image
	return result;
}

